import pixellib
from pixellib.tune_bg import alter_bg
import os
from utils.commonUtils import get_all_file_in_dir


def cut_bg(input_path, output_path, bg_color, detect):
    change_bg = alter_bg(model_type="pb")
    change_bg.load_pascalvoc_model("../utils/background_cut/models/xception_pascalvoc.pb")

    img_path_list = get_all_file_in_dir(input_path)

    # print(img_path_list)

    for i, img_path in enumerate(img_path_list):
        change_bg.color_bg(img_path, colors=bg_color, output_image_name=img_path.replace(input_path, output_path), detect=detect)
        print('cut_bg: {0}%'.format(round((i + 1) / len(img_path_list) * 100, 2)))








# class VideoSplit(object):
#     def __init__(self, weight_path):
#         self.change_bg = self.get_change_bg(weight_path)
#
#     def get_change_bg(self, weight_path):
#         '''
#         获取初始化分离模型
#         :return:
#         '''
#         change_bg = alter_bg(model_type="pb")
#         change_bg.load_pascalvoc_model(weight_path)
#         return change_bg
#
#     def split_video(self, video_name, bg_color, detect, frames_per_second=25):
#         # self.change_bg.change_video_bg('../video/inp/{0}'.format(video_name), "../img/bg/{0}.jpg".format(bg_color), frames_per_second=frames_per_second, output_video_name='../video/oup/{0}'.format(video_name), detect=detect)
#         self.change_bg.color_bg('../video/inp/{0}'.format(video_name), colors=(255, 255, 255), output_image_name="../video/oup/{0}".format(video_name), detect=detect)
#
#
# if __name__ == '__main__':
#
#     # props ----------------
#     video_name = "v111.mp4"
#     detect = "person"
#     bg_color = "white"
#     # ----------------------
#
#     vs = VideoSplit("../weights/xception_pascalvoc.pb")
#     vs.split_video(video_name=video_name, bg_color=bg_color, detect=detect)